#ifndef _THREADPOOL_H_
#define _THREADPOOL_H_

#include "stdafx.h"

// #include <windows.h>
#include <process.h> //_beginThreadEx
#include <stdio.h>

#include <intrin.h>

#include "exception.h"

static unsigned int pv_keepalive = 1;

typedef unsigned int (__stdcall *THREADPOOL_ENTRYPOINT)(volatile void* args);

#define MAX_JOBS_CNT 128

#define JOB_FINISHED 1
#define JOB_INIT 2
#define JOB_RUNNING 3

#define THREAD_SLEEPING 0
#define THREAD_RUNNING 1

class CThreadException : public CException
{
private:
	const char* _txt;
public:
	CThreadException(const char* reason):CException(reason){}
	~CThreadException(){}
};

struct sJob
{
	volatile struct sJob* next;
	volatile long state;
	volatile THREADPOOL_ENTRYPOINT fn;
	volatile void* arg;
	sJob():state(JOB_FINISHED), next(NULL){}
};

union joblist_head
{
	struct
	{
		volatile sJob* ptr;
		LONG magic;
	}st;
	volatile LONGLONG qw;
	joblist_head(LONGLONG val = 0):qw(val){}
};

struct sThreadInfo
{
	unsigned int idx;
	volatile long state;
	DWORD affMask;
	volatile joblist_head* jobs;
	HANDLE* threads;
	HANDLE* events;
	sThreadInfo(): threads(NULL), events(NULL), jobs(0), idx(0), state(THREAD_SLEEPING), affMask(1){}
};

class CHardwareInfo
{
private:
	DWORD _nb_cores;
	DWORD _process_mask;
public:
	CHardwareInfo()
	{
		DWORD sm;

		GetProcessAffinityMask(GetCurrentProcess(), &_process_mask, &sm);
		for (_nb_cores = 0; sm; _nb_cores++)
			sm &= sm - 1;
	}
	~CHardwareInfo(){}
	DWORD getNextThreadMask()
	{
		static unsigned char inc = 0;
		return 0x1 << (inc++ % _nb_cores);
	}
	DWORD getCoreNomber(void){return _nb_cores;}
};

static unsigned int __stdcall pv_thread_wrapper(void* args)
{
	// data initialization
	sThreadInfo* th_info = (sThreadInfo*) args;
	unsigned long id = th_info->idx;
	volatile joblist_head* head = th_info->jobs;

	SetThreadAffinityMask(th_info->threads[id], th_info->affMask);

	while (pv_keepalive)
	{
		volatile joblist_head oh;
		volatile joblist_head nh;

		// process jobs from list
		for (;;)
		{
			for (;;)
			{
				_InterlockedCompareExchange64(&oh.qw, head->qw, oh.qw);

				if (oh.st.ptr == NULL)
					goto end;

				nh.st.ptr = oh.st.ptr->next;
				nh.st.magic = oh.st.magic + 1;

				if (_InterlockedCompareExchange64(&head->qw, nh.qw ,oh.qw) == oh.qw)
					break;
			}

			volatile sJob* job = oh.st.ptr;

			job->fn(job->arg);
			InterlockedExchange(&(job->state), JOB_FINISHED);
		}
end:
		InterlockedExchange(&(th_info->state), THREAD_SLEEPING);

		// we use events to notify the main thread during a call to fence()
		SetEvent(th_info->events[id]);
		SuspendThread(th_info->threads[id]);
		ResetEvent(th_info->events[id]);
	}
	return 0;
}

class CThreadPool
{
private:
	sJob _jobs[MAX_JOBS_CNT];
	sThreadInfo* _th_info;
	joblist_head _joblist_head;
	HANDLE* _events;
	HANDLE* _threads;
	CHardwareInfo hwInfo;
	DWORD _nb_threads;
public:
	CThreadPool(int nb_threads = 0): _joblist_head(NULL), hwInfo()
	{
		unsigned int s;

		_nb_threads = (nb_threads) ? nb_threads : hwInfo.getCoreNomber();

		if (!(_threads = new HANDLE[nb_threads]))
			throw CThreadException("allocation error ! (threads)");
		if (!(_events = new HANDLE[nb_threads]))
		{
			delete _threads;
			throw CThreadException("allocation error ! (events)");
		}
		if (!(_th_info = new sThreadInfo[nb_threads]))
		{
			delete _threads;
			delete _events;
			throw CThreadException("allocation error ! (thread_info)");
		}

		for (unsigned char i = 0; i < nb_threads; i++)
		{
			_th_info[i].idx = i;
			_th_info[i].affMask = hwInfo.getNextThreadMask();
			_th_info[i].events = _events;
			_th_info[i].threads = _threads;
			_th_info[i].jobs = &_joblist_head;

			if ((_events[i] = CreateEvent(NULL, TRUE, FALSE, NULL)) == NULL)
			{
				delete _threads;
				delete _events;
				delete _th_info;
				throw CThreadException("CreateEvent() failed");
			}

			if ((_threads[i] = (HANDLE) _beginthreadex(NULL, 0, pv_thread_wrapper, (void *) &(_th_info[i]), CREATE_SUSPENDED, &s)) <= 0)
			{
				delete _threads;
				delete _events;
				delete _th_info;
				throw CThreadException("_beginthreadex() failed");
			}
		}
	}

	~CThreadPool()
	{
		pv_keepalive = 0;
	}

	BOOL postJob(volatile THREADPOOL_ENTRYPOINT pfn, volatile void* args)
	{
		// look for any free slot in the job list
		for (int i = 0; i < MAX_JOBS_CNT; i++)
		{
			if (InterlockedCompareExchange(&(_jobs[i].state), JOB_INIT, JOB_FINISHED) == JOB_FINISHED)
			{
				// free job slot found, creating object
				_jobs[i].fn = pfn;
				_jobs[i].arg = args;
				_jobs[i].next = NULL;

				joblist_head oh;
				joblist_head nh;

				// append new job to job stack, first position (higher priority)
				for (;;)
				{
					_InterlockedCompareExchange64(&(oh.qw), _joblist_head.qw, oh.qw);
					nh.st.magic = oh.st.magic + 1;
					nh.st.ptr = &(_jobs[i]);
					nh.st.ptr->next = oh.st.ptr;
					if (_InterlockedCompareExchange64(&_joblist_head.qw, nh.qw ,oh.qw) == oh.qw)
						break;
				}

				// try to wake-up a thread and nicely exit

				for (unsigned int j = 0; j < _nb_threads; j++)
				{
					if (InterlockedCompareExchange(&(_th_info[j].state), THREAD_RUNNING, THREAD_SLEEPING) == THREAD_SLEEPING)
					{
						ResumeThread(_threads[j]);
						break;
					}
				}
				return TRUE;
			}
		}
		// no free job slot found, performing a pool.fence() should be a wise move
		return FALSE;
	}

	void fence(void)
	{
		WaitForMultipleObjects(_nb_threads, _events, TRUE, INFINITE);
	}
};
#endif //_THREADPOOL_H_